package com.heima.datastructure.huffmantree;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 哈夫曼树 基于满二叉树（每个节点要么有两个孩子，要么是叶子节点）
 *
 * @author 勾新杰
 * @version 1.0
 * @date 2024/10/30 10:52
 */
public class HuffmanTree {
    private final Node root; // 哈夫曼树的根节点
    private final String data; // 存储传输数据的原始字符串
    private final Map<Character, Node> map; // 哈夫曼树节点存储的字符和对应的节点

    /**
     * 构造方法，用于构建哈夫曼树
     *
     * @param data 传输数据的原始字符串
     */
    @SuppressWarnings("all")
    public HuffmanTree(String data) {
        // 1. 对属性进行初始化
        this.data = data;
        this.map = new HashMap<>();
        // 2. 统计字符出现的频率
        char[] array = data.toCharArray();
        for (char c : array) {
            map.computeIfAbsent(c, Node::new);
            map.get(c).frequency++;
        }
        // 3. 构建哈夫曼树
        // 3.1 创建优先队列，用于存储节点
        PriorityQueue<Node> queue = new PriorityQueue<>(Comparator.comparingInt(Node::frequency));
        // 3.2 将节点加入优先队列
        queue.addAll(map.values());
        // 3.3 构建哈夫曼树，每次出队两个频率最小的节点，这两个作为子树组成一个新的子树，再重新入队
        while (queue.size() >= 2) {
            Node x = queue.poll();
            Node y = queue.poll();
            queue.offer(new Node(x.frequency + y.frequency, x, y));
        }
        // 3.4 初始化root
        this.root = queue.poll();
        // 4. 给每个叶子节点的编码code属性赋值，很明显，要用到深度优先遍历的思想
        encodePerChar(root, new StringBuilder());
    }

    /**
     * 对哈夫曼树的叶子节点进行编码存储
     * @param root 哈夫曼树根节点
     * @param builder builder
     */
    private void encodePerChar(Node root, StringBuilder builder) {
        // 递归结束条件：当前节点是叶子节点，结束时还要计算叶子节点的哈夫曼编码
        if (root.isLeaf()){
            root.code = builder.toString();
            builder.deleteCharAt(builder.length() - 1);
            return;
        }
        // 递的时候往左走则builder末尾新增一个0
        encodePerChar(root.left, builder.append(0));
        // 递的时候往右走则builder末尾新增一个1
        encodePerChar(root.right, builder.append(1));
        // 归的时候去掉builder最后一个0或1
        if (!builder.isEmpty()) { // 归到第一轮的时候
            builder.deleteCharAt(builder.length() - 1);
        }
    }

    /**
     * 哈夫曼树的节点类
     */
    static class Node {
        Character data; // 节点存储的字符，只有在叶子节点才有数据
        int frequency; // 节点存储的字符在传输数据中出现的频率
        Node left; // 左孩子
        Node right; // 右孩子
        String code; // 节点存储的字符对应的哈夫曼编码

        /**
         * 用于构建叶子节点
         *
         * @param data 叶子节点存储的字符
         */
        public Node(Character data) {
            this.data = data;
        }

        /**
         * 用于构建非叶子节点
         *
         * @param frequency 左右子树的叶子节点频率之和
         * @param left      左子树
         * @param right     右子树
         */
        public Node(int frequency, Node left, Node right) {
            this.frequency = frequency;
            this.left = left;
            this.right = right;
        }

        /**
         * 获取节点存储的字符在传输数据中出现的频率
         *
         * @return 节点存储的字符在传输数据中出现的频率
         */
        public int frequency() {
            return frequency;
        }

        /**
         * 判断是否是叶子节点
         *
         * @return true：是叶子节点，false：不是叶子节点
         */
        public boolean isLeaf() {
            return left == null; // 满二叉树，要么没有孩子，要么有两个孩子
        }

        /**
         * 重写toString方法，便于查看叶子节点信息
         *
         * @return 节点信息
         */
        @Override
        public String toString() {
            return "Node{" +
                    "data=" + data +
                    ", frequency=" + frequency +
                    ", code='" + code + '\'' +
                    '}';
        }
    }

    /**
     * 获取该字符串（传输数据）编码后的总比特位数
     * @return 总比特位数
     */
    public int codeLength(){
        int length = 0;
        for (Node node : map.values()) {
            length += node.frequency * node.code.length();
        }
        return length;
    }

    /**
     * 对传输数据进行编码
     * @return 编码后的结果，01这里用的字符串存储
     */
    public String encode(){
        char[] array = data.toCharArray();
        StringBuilder builder = new StringBuilder();
        for (char c : array) {
            builder.append(map.get(c).code);
        }
        return builder.toString();
    }

    /**
     * 对编码后的结果进行解码
     * @param code 编码后的结果
     * @return 解码结果
     */
    public String decode(String code){
        Node curr = root;
        char[] array = code.toCharArray();
        StringBuilder builder = new StringBuilder();
        for (char c : array) {
            if (c == '0'){
                curr = curr.left;
            }else {
                curr = curr.right;
            }
            // 编码结果非空说明肯定有01所以肯定会向左或右走，所以root绝不可能是叶子节点，可以大胆的先移动一次curr
            // 而且编码结果的最后一次0或1移动后肯定是叶子节点，这里可以处理得到
            if (curr.isLeaf()){
                builder.append(curr.data);
                curr = root;
            }
        }
        return builder.toString();
    }
}
